"""
股票分析智能体模块 - 提供股票分析功能
"""

from typing import Dict, TypedDict, List, Union, Any, Callable
from langgraph.graph import StateGraph, END
import json
import datetime
import logging
import pandas as pd

from FDAS.agents.base import BaseAgent, BaseAgentState
from FDAS.data import get_asset
from FDAS.data.assets.base import ASSET_TYPE_STOCK
from FDAS.agents.workflows.prompts import (
    get_query_parse_prompt,
    get_specific_analysis_prompt,
    get_synthesis_prompt
)

# 自定义JSON编码器，处理日期类型
class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (datetime.date, datetime.datetime)):
            return obj.isoformat()
        return super(DateTimeEncoder, self).default(obj)

# 定义股票智能体状态
class StockAgentState(BaseAgentState):
    symbol: str = ""
    analysis_type: str = ""
    stock_data: Dict = {}
    # 各子分析结果
    basic_info_analysis: str = ""
    valuation_analysis: str = ""
    financial_analysis: str = ""
    historical_analysis: str = ""
    fund_flow_analysis: str = ""
    news_analysis: str = ""
    technical_analysis: str = ""

class StockAnalysisAgent(BaseAgent):
    """股票分析智能体"""
    
    def _build_workflow(self) -> StateGraph:
        """构建智能体工作流"""
        workflow = StateGraph(StockAgentState)
        
        # 添加节点
        workflow.add_node("parse_query", self._parse_query)
        workflow.add_node("retrieve_data", self._retrieve_data)
        workflow.add_node("analyze_data", self._analyze_data)
        
        # 设置节点连接关系
        workflow.set_entry_point("parse_query")
        workflow.add_edge("parse_query", "retrieve_data")
        workflow.add_edge("retrieve_data", "analyze_data")
        workflow.add_edge("analyze_data", END)
        
        return workflow
    
    def _parse_query(self, state: StockAgentState) -> Dict:
        """解析用户查询并提取股票代码和分析类型"""
        try:
            response = self.llm.invoke(
                "从以下股票分析请求中提取信息：\n" + state['user_query'] +
                "\n只输出JSON格式，不要任何解释，格式为：\n" +
                '{"symbol": "股票代码", "analysis_type": "基本面分析|估值分析|财务分析|历史分析|资金流分析|新闻分析|技术分析|综合分析"}\n' +
                '确保：1.使用双引号 2.无注释 3.代码通常为6位数字'
            )
            
            # 解析JSON
            result = self._clean_json_response(response.content)
            
            # 数据校验
            result['symbol'] = str(result.get('symbol', '000001')).strip().zfill(6)
            analysis_type = str(result.get('analysis_type', '基本面分析')).lower()
            result['analysis_type'] = {
                'basic': '基本面分析',
                'valuation': '估值分析',
                'financial': '财务分析',
                'historical': '历史分析',
                'fund_flow': '资金流分析',
                'news': '新闻分析',
                'technical': '技术分析',
                'comprehensive': '综合分析',
                '基本面': '基本面分析',
                '估值': '估值分析',
                '财务': '财务分析',
                '历史': '历史分析',
                '资金流': '资金流分析',
                '新闻': '新闻分析',
                '技术': '技术分析',
                '综合': '综合分析'
            }.get(analysis_type, '基本面分析')
            
            return {
                "symbol": result['symbol'],
                "analysis_type": result['analysis_type'],
                "parsed_query": json.dumps(result, ensure_ascii=False, cls=DateTimeEncoder),
                "user_query": state['user_query']
            }
                
        except Exception as e:
            self.logger.error(f"查询解析失败: {str(e)}")
            # 默认返回
            return {
                "symbol": "000001",
                "analysis_type": "基本面分析",
                "parsed_query": f"解析失败：{str(e)}",
                "user_query": state['user_query']
            }
    
    def _retrieve_data(self, state: StockAgentState) -> Dict:
        """获取相关数据"""
        try:
            symbol = state['symbol']
            analysis_type = state['analysis_type']
            
            # 获取股票对象
            stock = get_asset(symbol, ASSET_TYPE_STOCK)
            
            # 基础数据 - 所有分析类型都需要
            basic_info = stock.basic_info
            
            # 根据分析类型获取不同数据
            if analysis_type == '基本面分析':
                data = {
                    "basic_info": basic_info,
                    "industry": stock.get_industry_info(),
                    "shareholders": stock.get_shareholders(),
                    "history": stock.get_history(period="daily", limit=30).to_dict('records')
                }
            elif analysis_type == '估值分析':
                data = {
                    "basic_info": basic_info,
                    "valuation": stock.get_valuation(),
                    "history": stock.get_history(period="daily", limit=60).to_dict('records')
                }
            elif analysis_type == '财务分析':
                data = {
                    "basic_info": basic_info,
                    "financial": {
                        "balance_sheet": stock.get_balance_sheet(limit=2).to_dict('records'),
                        "income_statement": stock.get_income_statement(limit=2).to_dict('records'),
                        "cash_flow": stock.get_cash_flow(limit=2).to_dict('records'),
                        "indicators": stock.get_financial_indicators(limit=2).to_dict('records')
                    },
                    "history": stock.get_history(period="daily", limit=30).to_dict('records')
                }
            elif analysis_type == '历史分析':
                data = {
                    "basic_info": basic_info,
                    "history": stock.get_history(period="daily", limit=180).to_dict('records'),
                    "dividends": stock.get_dividends()
                }
            elif analysis_type == '资金流分析':
                data = {
                    "basic_info": basic_info,
                    "fund_flow": stock.get_fund_flow(limit=30),
                    "history": stock.get_history(period="daily", limit=30).to_dict('records')
                }
            elif analysis_type == '新闻分析':
                data = {
                    "basic_info": basic_info,
                    "news": stock.get_news(limit=10),
                    "history": stock.get_history(period="daily", limit=30).to_dict('records')
                }
            elif analysis_type == '技术分析':
                data = {
                    "basic_info": basic_info,
                    "technical": stock.get_technical_indicators(limit=30),
                    "history": stock.get_history(period="daily", limit=60).to_dict('records')
                }
            else:  # 综合分析
                data = {
                    "basic_info": basic_info,
                    "industry": stock.get_industry_info(),
                    "valuation": stock.get_valuation(),
                    "financial": {
                        "balance_sheet": stock.get_balance_sheet(limit=2).to_dict('records'),
                        "income_statement": stock.get_income_statement(limit=2).to_dict('records'),
                        "cash_flow": stock.get_cash_flow(limit=2).to_dict('records'),
                        "indicators": stock.get_financial_indicators(limit=2).to_dict('records')
                    },
                    "history": stock.get_history(period="daily", limit=90).to_dict('records'),
                    "fund_flow": stock.get_fund_flow(limit=30),
                    "news": stock.get_news(limit=5),
                    "technical": stock.get_technical_indicators(limit=30)
                }
            
            return {
                **state,  # 保留现有状态
                "stock_data": data
            }
                
        except Exception as e:
            self.logger.error(f"数据获取失败: {str(e)}")
            return {
                **state,
                "stock_data": {"error": f"数据获取失败: {str(e)}"}
            }
    
    def _analyze_data(self, state: StockAgentState) -> Dict:
        """执行分析并生成报告"""
        try:
            # 检查数据
            stock_data = state.get('stock_data', {})
            
            if not stock_data:
                return {**state, "analysis_result": self._format_error_response("未获取到有效数据")}
                
            if 'error' in stock_data:
                return {**state, "analysis_result": self._format_error_response(stock_data['error'])}
            
            # 构建分析提示
            analysis_type = state['analysis_type']
            
            # 根据分析类型构建不同的提示
            if analysis_type == '基本面分析':
                prompt = f"""
                请基于以下数据对股票进行基本面分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                行业信息：{json.dumps(stock_data.get('industry', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                股东信息：{json.dumps(stock_data.get('shareholders', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票基本面分析报告，包含以下要素：
                1. 公司概况
                2. 行业地位分析
                3. 股东结构分析
                4. 近期股价表现
                5. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '估值分析':
                prompt = f"""
                请基于以下数据对股票进行估值分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                估值数据：{json.dumps(stock_data.get('valuation', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票估值分析报告，包含以下要素：
                1. 当前估值水平
                2. 历史估值比较
                3. 同行业估值比较
                4. 估值合理性分析
                5. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '财务分析':
                prompt = f"""
                请基于以下数据对股票进行财务分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                财务数据：{json.dumps(stock_data.get('financial', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票财务分析报告，包含以下要素：
                1. 资产负债表分析
                2. 利润表分析
                3. 现金流量表分析
                4. 财务指标分析
                5. 财务健康状况评估
                6. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '历史分析':
                prompt = f"""
                请基于以下数据对股票进行历史分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:10], ensure_ascii=False, cls=DateTimeEncoder)}
                
                分红数据：{json.dumps(stock_data.get('dividends', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票历史分析报告，包含以下要素：
                1. 长期价格趋势
                2. 波动性分析
                3. 关键支撑位和阻力位
                4. 成交量分析
                5. 分红历史分析
                6. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '资金流分析':
                prompt = f"""
                请基于以下数据对股票进行资金流分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                资金流数据：{json.dumps(stock_data.get('fund_flow', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票资金流分析报告，包含以下要素：
                1. 主力资金流向
                2. 散户资金流向
                3. 资金流与价格关系
                4. 资金流异常分析
                5. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '新闻分析':
                prompt = f"""
                请基于以下数据对股票进行新闻分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                新闻数据：{json.dumps(stock_data.get('news', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票新闻分析报告，包含以下要素：
                1. 新闻情感分析
                2. 重大事件影响
                3. 新闻与价格关系
                4. 舆情风险分析
                5. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '技术分析':
                prompt = f"""
                请基于以下数据对股票进行技术分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                技术指标：{json.dumps(stock_data.get('technical', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:10], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的股票技术分析报告，包含以下要素：
                1. 趋势分析
                2. 支撑与阻力位
                3. 移动平均线分析
                4. MACD、KDJ等指标分析
                5. 成交量分析
                6. 技术形态分析
                7. 交易建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            else:  # 综合分析
                prompt = f"""
                请基于以下数据对股票进行综合分析：
                
                基本信息：{json.dumps(stock_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                行业信息：{json.dumps(stock_data.get('industry', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                估值数据：{json.dumps(stock_data.get('valuation', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                财务数据：{json.dumps(stock_data.get('financial', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(stock_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                资金流数据：{json.dumps(stock_data.get('fund_flow', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                新闻数据：{json.dumps(stock_data.get('news', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                技术指标：{json.dumps(stock_data.get('technical', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份全面的股票综合分析报告，包含以下要素：
                1. 核心结论摘要
                2. 公司基本面分析
                3. 估值分析
                4. 财务分析
                5. 技术面分析
                6. 资金流向分析
                7. 新闻舆情分析
                8. 风险因素
                9. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            
            # 执行分析
            analysis = self.llm.invoke(prompt)
            
            return {**state, "analysis_result": analysis.content}
                
        except Exception as e:
            self.logger.error(f"分析过程中发生错误: {str(e)}")
            return {**state, "analysis_result": self._format_error_response(str(e))}
    
    def analyze(self, query: str) -> str:
        """执行完整分析流程"""
        try:
            result = self.agent.invoke({"user_query": query})
            return result['analysis_result']
        except Exception as e:
            self.logger.error(f"分析执行失败: {str(e)}")
            return self._format_error_response(str(e))
    
    def _analyze_financial(self, state: StockAgentState) -> Dict:
        """分析财务数据"""
        try:
            # 获取相关数据
            stock_data = state.get('stock_data', {})
            basic_info = stock_data.get('basic_info', {})
            financial_data = stock_data.get('financial', {})
            
            # 检查数据
            if not financial_data:
                return {**state, "financial_analysis": "未获取到财务数据"}
            
            # 限制财务数据量 - 只保留最近2期
            simplified_financial = {}
            
            # 对资产负债表进行精简 - 只保留最近2期，且只保留关键指标
            if 'balance_sheet' in financial_data and financial_data['balance_sheet']:
                # 取前2期
                raw_balance = financial_data['balance_sheet'][:2]
                # 只保留关键指标
                key_balance_metrics = ['总资产', '总负债', '股东权益合计', '货币资金', '应收账款', '存货', '流动资产合计', '流动负债合计']
                simplified_balance = []
                for item in raw_balance:
                    simplified_item = {k: item.get(k) for k in key_balance_metrics if k in item}
                    simplified_item['报告期'] = item.get('报告期', '')
                    simplified_balance.append(simplified_item)
                simplified_financial['balance_sheet'] = simplified_balance
            else:
                simplified_financial['balance_sheet'] = []
                
            # 对利润表进行精简 - 只保留最近2期，且只保留关键指标
            if 'income_statement' in financial_data and financial_data['income_statement']:
                # 取前2期
                raw_income = financial_data['income_statement'][:2]
                # 只保留关键指标
                key_income_metrics = ['营业收入', '营业利润', '利润总额', '净利润', '归属于母公司股东的净利润', '毛利率', '净利率']
                simplified_income = []
                for item in raw_income:
                    simplified_item = {k: item.get(k) for k in key_income_metrics if k in item}
                    simplified_item['报告期'] = item.get('报告期', '')
                    simplified_income.append(simplified_item)
                simplified_financial['income_statement'] = simplified_income
            else:
                simplified_financial['income_statement'] = []
                
            # 对现金流量表进行精简 - 只保留最近2期，且只保留关键指标
            if 'cash_flow' in financial_data and financial_data['cash_flow']:
                # 取前2期
                raw_cash = financial_data['cash_flow'][:2]
                # 只保留关键指标
                key_cash_metrics = ['经营活动产生的现金流量净额', '投资活动产生的现金流量净额', '筹资活动产生的现金流量净额', '现金及现金等价物净增加额']
                simplified_cash = []
                for item in raw_cash:
                    simplified_item = {k: item.get(k) for k in key_cash_metrics if k in item}
                    simplified_item['报告期'] = item.get('报告期', '')
                    simplified_cash.append(simplified_item)
                simplified_financial['cash_flow'] = simplified_cash
            else:
                simplified_financial['cash_flow'] = []
                
            # 对财务指标进行精简 - 只保留最近2期，且只保留关键指标
            if 'indicators' in financial_data and financial_data['indicators']:
                # 取前2期
                raw_indicators = financial_data['indicators'][:2]
                # 只保留关键指标
                key_indicator_metrics = ['ROE', 'ROA', '资产负债率', '流动比率', '速动比率', '每股收益', '每股净资产']
                simplified_indicators = []
                for item in raw_indicators:
                    simplified_item = {k: item.get(k) for k in key_indicator_metrics if k in item}
                    simplified_item['报告期'] = item.get('报告期', '')
                    simplified_indicators.append(simplified_item)
                simplified_financial['indicators'] = simplified_indicators
            else:
                simplified_financial['indicators'] = []
            
            # 构建提示
            prompt = get_specific_analysis_prompt('financial', 'stock').format(
                basic_info=json.dumps(basic_info, ensure_ascii=False, cls=DateTimeEncoder),
                financial_data=json.dumps(simplified_financial, ensure_ascii=False, cls=DateTimeEncoder)
            )
            
            # 调用大模型
            response = self.llm.invoke(prompt)
            
            # 返回结果
            return {**state, "financial_analysis": response.content}
                
        except Exception as e:
            self.logger.error(f"财务分析失败: {str(e)}")
            return {**state, "financial_analysis": f"分析失败: {str(e)}"}